#include "ConstPropagation.hpp"
#include "logging.hpp"
#include "LoopSearch.hpp"

// 给出了返回整形值的常数折叠实现，大家可以参考，在此基础上拓展
// 当然如果同学们有更好的方式，不强求使用下面这种方式
ConstantInt *ConstFolder::compute(
    Instruction::OpID op,
    ConstantInt *value1,
    ConstantInt *value2)
{

    int c_value1 = value1->get_value();
    int c_value2 = value2->get_value();
    switch (op)
    {
    case Instruction::add:
        return ConstantInt::get(c_value1 + c_value2, module_);
        break;
    case Instruction::sub:
        return ConstantInt::get(c_value1 - c_value2, module_);
        break;
    case Instruction::mul:
        return ConstantInt::get(c_value1 * c_value2, module_);
        break;
    case Instruction::sdiv:
        return ConstantInt::get((int)(c_value1 / c_value2), module_);
        break;
    default:
        return nullptr;
        break;
    }
}
ConstantFP *ConstFolder::compute(
    Instruction::OpID op,
    ConstantFP *value1,
    ConstantFP *value2)
{

    float c_value1 = value1->get_value();
    float c_value2 = value2->get_value();
    switch (op)
    {
    case Instruction::fadd:
        return ConstantFP::get(c_value1 + c_value2, module_);
        break;
    case Instruction::fsub:
        return ConstantFP::get(c_value1 - c_value2, module_);
        break;
    case Instruction::fmul:
        return ConstantFP::get(c_value1 * c_value2, module_);
        break;
    case Instruction::fdiv:
        return ConstantFP::get(c_value1 / c_value2, module_);
        break;
    default:
        return nullptr;
        break;
    }
}

Constant *ConstFolder::compute(
    CmpInst::CmpOp op,
    ConstantInt *value1,
    ConstantInt *value2)
{
    int c_value1 = value1->get_value();
    int c_value2 = value2->get_value();

    switch (op)
    {
    case CmpInst::CmpOp::EQ:
        return ConstantInt::get(c_value1 == c_value2, module_);
        break;
    case CmpInst::CmpOp::NE:
        return ConstantInt::get(c_value1 != c_value2, module_);
        break;
    case CmpInst::CmpOp::LE:
        return ConstantInt::get(c_value1 <= c_value2, module_);
        break;
    case CmpInst::CmpOp::LT:
        return ConstantInt::get(c_value1 < c_value2, module_);
        break;
    case CmpInst::CmpOp::GE:
        return ConstantInt::get(c_value1 >= c_value2, module_);
        break;
    case CmpInst::CmpOp::GT:
        return ConstantInt::get(c_value1 > c_value2, module_);
        break;
    default:
        return nullptr;
    }
}

Constant *ConstFolder::compute(
    FCmpInst::CmpOp op,
    ConstantFP *value1,
    ConstantFP *value2)
{
    float c_value1 = value1->get_value();
    float c_value2 = value2->get_value();

    switch (op)
    {
    case FCmpInst::CmpOp::EQ:
        return ConstantInt::get(c_value1 == c_value2, module_);
        break;
    case FCmpInst::CmpOp::NE:
        return ConstantInt::get(c_value1 != c_value2, module_);
        break;
    case FCmpInst::CmpOp::LE:
        return ConstantInt::get(c_value1 <= c_value2, module_);
        break;
    case FCmpInst::CmpOp::LT:
        return ConstantInt::get(c_value1 < c_value2, module_);
        break;
    case FCmpInst::CmpOp::GE:
        return ConstantInt::get(c_value1 >= c_value2, module_);
        break;
    case FCmpInst::CmpOp::GT:
        return ConstantInt::get(c_value1 > c_value2, module_);
        break;
    default:
        return nullptr;
    }
}

Constant *ConstFolder::compute(
    Instruction::OpID op,
    Constant *value)
{
    auto ivalue = dynamic_cast<ConstantInt *>(value);
    auto fvalue = dynamic_cast<ConstantFP *>(value);

    switch (op)
    {
    case Instruction::fptosi:
        return ConstantInt::get(int(fvalue->get_value()), module_);
    case Instruction::sitofp:
        return ConstantFP::get(float(ivalue->get_value()), module_);
    case Instruction::zext:
        return ConstantInt::get(bool(ivalue->get_value()), module_);
    default:
        return nullptr;
    }
}

Constant *ConstFolder::compute(Instruction *ins, std::vector<Constant *> args)
{
    if (ins->is_add() || ins->is_sub() || ins->is_mul() || ins->is_div())
    {
        auto value1 = dynamic_cast<ConstantInt *>(args[0]);
        auto value2 = dynamic_cast<ConstantInt *>(args[1]);
        return compute(ins->get_instr_type(), value1, value2);
    }
    else if (ins->is_fadd() || ins->is_fsub() || ins->is_fmul() || ins->is_fdiv())
    {
        auto value1 = dynamic_cast<ConstantFP *>(args[0]);
        auto value2 = dynamic_cast<ConstantFP *>(args[1]);
        return compute(ins->get_instr_type(), value1, value2);
    }
    else if (ins->is_fp2si() || ins->is_si2fp() || ins->is_zext())
    {
        return compute(ins->get_instr_type(), args[0]);
    }
    else if (ins->is_cmp())
    {
        auto cmp_ins = dynamic_cast<CmpInst *>(ins);
        auto op = cmp_ins->get_cmp_op();
        auto value1 = dynamic_cast<ConstantInt *>(args[0]);
        auto value2 = dynamic_cast<ConstantInt *>(args[1]);
        return compute(op, value1, value2);
    }
    else if (ins->is_fcmp())
    {
        auto fcmp_ins = dynamic_cast<FCmpInst *>(ins);
        auto op = fcmp_ins->get_cmp_op();
        auto value1 = dynamic_cast<ConstantFP *>(args[0]);
        auto value2 = dynamic_cast<ConstantFP *>(args[1]);
        return compute(op, value1, value2);
    }
    return nullptr;
}


/**
 * @brief 递归地执行常量替换和死代码删除，执行完一个基本块后进入其后继块
 * 
 * @param bb 当前进入的基本块
 * @param var_need_replace， 用于存放需要被替换成常量的变量
 * @param bb_entried 存放已经进入过的基本块
 */
void ConstPropagation::ConstPropagate(BasicBlock *bb, std::map<std::string, Constant *> var_need_replace, std::set<BasicBlock*> *bb_entried){
    bb_entried->insert(bb);
    std::vector<Instruction *> ins2del;
    for (auto instruction: bb->get_instructions())
    {
        //常量替换
        std::vector<Constant*> constant_args;
        //将操作数中的变量替换为常量
        for (size_t i = 0; i < instruction->get_num_operand(); i++)
        {
            auto val = instruction->get_operand(i);
            constant_args.push_back(dynamic_cast<ConstantInt*> (val));
            if (constant_args.back() == nullptr)
            {
                constant_args[i] = dynamic_cast<ConstantFP*> (val);
            }
            if (var_need_replace.count(val->get_name()) != 0)
            {
                instruction->set_operand(i, var_need_replace[val->get_name()]);
                constant_args[i] = var_need_replace[val->get_name()];
            }
        }

        //将左值计算成绝对值，替换后续用到的左值，删除计算式
        bool should_replace = true; //左值有可能被完全替换成常量
        for (auto val: constant_args)
        {
            if (val == nullptr)
            {
                should_replace = false;
                break;
            }
        }

        if (should_replace&&(ConstFolder(this->m_).compute(instruction, constant_args)!=nullptr))
        {
            auto left_val = ConstFolder(this->m_).compute(instruction, constant_args);
                instruction->remove_use_of_ops();
                instruction->replace_all_use_with(left_val);
                ins2del.push_back(instruction);
                var_need_replace.insert({instruction->get_name(), left_val});
        }
        //删掉phi指令中冗余变量
        if(instruction->is_phi())
        {
            std::set<BasicBlock *> prebb;
            for(auto pre: bb->get_pre_basic_blocks())
            {
                prebb.insert(pre);
            }
            for(auto i=0;i<instruction->get_num_operand();i+=2)
            {
                auto phibb = dynamic_cast<BasicBlock *>(instruction->get_operand(i +1));
                if(!prebb.count(phibb))
                {
                    instruction->remove_operands(i,i+1);
                }


            }
        }



        //常量替换完成，执行死代码删除
        //常量替换时已经删除了被完全替换成常量的而产生的死代码
        //但还需要删除因为常数条件跳转而永远无法到达的基本块
        //遇到条件跳转指令时处理
        if (instruction->is_br())
        {
            auto instruction_br = dynamic_cast<BranchInst*> (instruction);
            if (instruction_br->is_cond_br())
            {
                //跳转条件
                auto condition = dynamic_cast<ConstantInt*> (instruction_br->get_operand(0));
                if (condition != nullptr || var_need_replace.count(instruction_br->get_operand(0)->get_name()) != 0)
                {
                    condition = dynamic_cast<ConstantInt*> (var_need_replace[instruction_br->get_operand(0)->get_name()]);
                }
                if(condition != nullptr){
                    auto succ_bb = dynamic_cast<BasicBlock *>(instruction_br->get_operand(condition->get_value() ? 1: 2));
                    auto delete_bb = dynamic_cast<BasicBlock *>(instruction_br->get_operand(condition->get_value() ? 2: 1));
                    instruction_br->remove_operands(0, 2);
                    instruction_br->add_operand(succ_bb);
                    bb->remove_succ_basic_block(delete_bb);
                }
            }
        }
    }
    for (auto ins : ins2del)
    {
        bb->delete_instr(ins);
    }
    for (auto succ_bb: bb->get_succ_basic_blocks())
    {
        if (bb_entried->count(succ_bb) == 0)
        {
            ConstPropagate(succ_bb, var_need_replace, bb_entried);
        }
        
    }
    
    
}

void ConstPropagation::run()
{
    // 从这里开始吧！
    auto M = ConstFolder(m_);
    for (auto func : m_->get_functions())
    {
        if (func->get_basic_blocks().size() == 0)
        {
            continue;
        }
        std::map<std::string, Constant *> var_need_replace;
        std::set<BasicBlock *> bb_entried;
        auto bb = func->get_entry_block();
        ConstPropagate(bb, var_need_replace, &bb_entried);
    }
}
